# -*- coding: utf-8 -*-
"""
Created on Fri Nov 19 12:12:54 2021

@author: zhiyinan
"""

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import Parameters_1D as par_RNN 
from Polar1D_vectorized import *
import van_Genuchten1D_vectorized as vg 
from scipy import integrate
from sklearn.preprocessing import MinMaxScaler
from casadi import *
import matplotlib.font_manager as font_manager
import time
import pandas as pd
from sklearn.linear_model import Ridge
import random

# %% Define NN layers
def sigmoid(x):
    return 1/(1+np.exp(-x))

def LSTMlayer_sig(warr, uarr, barr, x_t,h_tm1,c_tm1):
    
    '''
    c_tm1 = np.array([0,0]).reshape(1,2)
    h_tm1 = np.array([0,0]).reshape(1,2)
    x_t   = np.array([1]).reshape(1,1)
    
    warr.shape = (nfeature,hunits*4)
    uarr.shape = (hunits,hunits*4)
    barr.shape = (hunits*4,)
    
    '''
    
    s_t = (mtimes(x_t,warr) + mtimes(h_tm1,uarr) + barr.reshape((1, barr.shape[0])))
    hunit = uarr.shape[0]
    i  = sigmoid(s_t[:,:hunit])
    f  = sigmoid(s_t[:,1*hunit:2*hunit])
    # _c = np.tanh(s_t[:,2*hunit:3*hunit])
    _c = sigmoid(s_t[:,2*hunit:3*hunit])
    o  = sigmoid(s_t[:,3*hunit:])
    c_t = i*_c + f*c_tm1
    # h_t = o*np.tanh(c_t)
    h_t = o*sigmoid(c_t)
    return h_t, c_t

def LSTMlayer_tanh(warr, uarr, barr, x_t,h_tm1,c_tm1):
    
    '''
    c_tm1 = np.array([0,0]).reshape(1,2)
    h_tm1 = np.array([0,0]).reshape(1,2)
    x_t   = np.array([1]).reshape(1,1)
    
    warr.shape = (nfeature,hunits*4)
    uarr.shape = (hunits,hunits*4)
    barr.shape = (hunits*4,)
    
    '''
    
    s_t = (mtimes(x_t,warr) + mtimes(h_tm1,uarr) + barr.reshape((1, barr.shape[0])))
    hunit = uarr.shape[0]
    i  = sigmoid(s_t[:,:hunit])
    f  = sigmoid(s_t[:,1*hunit:2*hunit])
    _c = np.tanh(s_t[:,2*hunit:3*hunit])
    # _c = sigmoid(s_t[:,2*hunit:3*hunit])
    o  = sigmoid(s_t[:,3*hunit:])
    c_t = i*_c + f*c_tm1
    h_t = o*np.tanh(c_t)
    # h_t = o*sigmoid(c_t)
    return h_t, c_t




# %% Define sub model 3

model_3 = tf.keras.models.load_model("Modeling/p20_f1_5s_3s_sig_2h_029_038_4_35_e7_n")
# model_3 = model_1
def M3(x):
    weightLSTM_1 = model_3.layers[0].get_weights()
    warr1, uarr1, barr1 = weightLSTM_1

    weightLSTM_2 = model_3.layers[1].get_weights()
    warr2, uarr2, barr2 = weightLSTM_2
    
    weightD = model_3.layers[2].get_weights()
    kernel3, bias3 = weightD
    
    n1 = 5
    # Initialise the model with the hidden states
    h_t1 = SX.zeros((1, n1))
    # Long term memory
    C_t1 = SX.zeros((1, n1))
    for i in range(x.shape[0]):
        xi = x[i,:].reshape((1,2))
        L1 = LSTMlayer_sig(warr1, uarr1, barr1, xi, h_t1[i,:], C_t1[i,:])
        h_t1 = vertcat(h_t1, L1[0])
        C_t1 = vertcat(C_t1, L1[1])
    
    n2 = 3
    # Initialise the model with the hidden states
    h_t2 = SX.zeros((1, n2))
    # Long term memory
    C_t2 = SX.zeros((1, n2))
    for i in range(h_t1.shape[0]):
        xi = h_t1[i,:].reshape((1,n1))
        L2 = LSTMlayer_sig(warr2, uarr2, barr2, xi, h_t2[i,:], C_t2[i,:])
        h_t2 = vertcat(h_t2, L2[0])
        C_t2 = vertcat(C_t2, L2[1])
    #  array([-2.7764497, -2.6578658, -2.6829786, -2.7356846, -2.73696  ]
    # Dense layer
    out2 = []
    for i in range(bias3.shape[0]):
        temp2 = mtimes(h_t2[-1,:], kernel3)[i]
    temp2 = bias3.reshape((1, bias3.shape[0])) + mtimes(h_t2[-1,:], kernel3)
    out2 = sigmoid(temp2)
    
    return out2


# def M3(x):
#     weightLSTM_1 = model_3.layers[0].get_weights()
#     warr1, uarr1, barr1 = weightLSTM_1

#     weightD = model_3.layers[1].get_weights()
#     kernel2, bias2 = weightD
    
#     n1 = 10
#     # Initialise the model with the hidden states
#     h_t1 = SX.zeros((1, n1))
#     # Long term memory
#     C_t1 = SX.zeros((1, n1))
#     for i in range(x.shape[0]):
#         xi = x[i,:].reshape((1,2))
#         L1 = LSTMlayer_sig(warr1, uarr1, barr1, xi, h_t1[i,:], C_t1[i,:])
#         h_t1 = vertcat(h_t1, L1[0])
#         C_t1 = vertcat(C_t1, L1[1])
    
#     #  array([-2.7764497, -2.6578658, -2.6829786, -2.7356846, -2.73696  ]
#     # Dense layer
#     out2 = []
#     for i in range(bias2.shape[0]):
#         temp2 = mtimes(h_t1[-1,:], kernel2)[i]
#     temp2 = bias2.reshape((1, bias2.shape[0])) + mtimes(h_t1[-1,:], kernel2)
#     out2 = sigmoid(temp2)     
    
#     return out2

# %% Input Signal generation
np.random.seed(3)
random.seed(3)


soilPars = loamySoil() # The soil parameters
totalDepth, axialNodes, axialDistance, totalNodes, numOfActuators, interPoints = par_RNN.spatialVariables_1D()  #Spatial parameters
samplingTime, samplingTimeInternal, internalTimeSteps, totTimeSteps = \
                                 par_RNN.temporalVariable_dataGen(1000*2*60*60) #time parameters

# lbu = -5.10e-05  #  unit [m/s]
# ubu = -4.94e-05  #  unit [m/s]

# irr_freq = int(20*60*60/samplingTime)   # PRBS change once only if a minimum of 100 hr is reached

uv  = np.zeros((totTimeSteps,1))


i=0
while i <= 1000:
    irr_freq = int(30*60*60/samplingTime)
    ulist = np.array([-0.45e-6, -2e-6, -1e-6, -1.5e-6, -3e-6, -2.5e-6 , -3.5e-6])
    # ulist = [-2/1000/3600, -1.5/1000/3600, -0.9/1000/3600, -1.2/1000/3600, -0.5/1000/3600]
    t = irr_freq + random.randrange(irr_freq)
    uv[i:i+t] = random.choice(ulist)
    i = i+t
    print(i)
    

plt.figure(1)
plt.plot(uv[:])                


initCondition = -0.1*np.ones(totalNodes)
#Arrays to store the results




def ODE(t, head, irrigAmount, cropCoeff, refEvap, soilPars):
    return Richards1D_vectorized(head, irrigAmount, cropCoeff, refEvap, soilPars)

cropCoeff = 0.88*np.ones(totTimeSteps)
refEvap = 3.102e-08*np.ones(totTimeSteps)

# %% Integration 

headArray= np.zeros((totTimeSteps+1, totalNodes)) # Pressure head, the state
headArray[0,:] = initCondition

volMoisture = np.zeros((totTimeSteps+1, totalNodes)) # Volumetric moisture, the measured output
volMoisture[0,:] = volMoistureAllNodes_1D(initCondition, soilPars).ravel()


i=1
temp = headArray[i-1]
ui = uv[i-1]
for t in range(int(samplingTime/samplingTimeInternal)): 
    sol = integrate.solve_ivp(ODE, args = (ui, cropCoeff[i-1], refEvap[i-1], soilPars),
                              t_span=[0,samplingTimeInternal],y0=tuple(temp),method='BDF')
    temp = sol.y[:,-1]
    # ui = 0
headArray[i,:]=sol.y[:,-1]
volMoisture[i,:] = volMoistureAllNodes_1D(headArray[i,:], soilPars).ravel()

for i in range(2, totTimeSteps + 1):  # totTimeSteps + 1
    print(i)
    temp = headArray[i-1]
    ui = uv[i-1]
    for t in range(int(samplingTime/samplingTimeInternal)): 
        sol = integrate.solve_ivp(ODE, args = (ui, cropCoeff[i-1], refEvap[i-1], soilPars),
                              t_span=[0,samplingTimeInternal],y0=tuple(temp),method='BDF')
        temp = sol.y[:,-1]
    # print(sol.y[:,-1])
        # ui = 0
    headArray[i,:]=sol.y[:,-1]
    volMoisture[i,:] = volMoistureAllNodes_1D(headArray[i,:], soilPars).ravel()

plt.figure(2)
plt.plot(volMoisture[1000:2000,19])

# np.savetxt("Modeling/u_train_30hrIF_30000_029_038.txt", uv)
# np.savetxt("Modeling/y_train_30hrIF_30000_029_038.txt", volMoisture)

# %%
past = 20
future = 1
y_index = 19

def multivariate_data(dataset, target, start_index, end_index, history_size,
                      target_size,u_shape): #, single_step=False):
    data = []
    labels = []

    start_index = start_index + history_size
    if end_index is None:
        end_index = len(dataset) - target_size
        
    for i in range(start_index, end_index):
        temp_p = dataset[i-history_size:i,:]
        for j in range(target_size-1):
        # temp_p = np.concatenate(((dataset[i-history_size:i,:]).flatten(), 
        #                            (dataset[i:i+target_size-1,:u_shape]).flatten()))
            temp = np.array([dataset[i+j,0], dataset[i,1]]).reshape((1,2))
            temp_p = np.concatenate((temp_p,temp))
        data.append(temp_p)
        labels.append((target[i:i+target_size].flatten()))
    

    return np.array(data), np.array(labels)


# %% Training data

u = np.loadtxt("Modeling/u_train_30hrIF_30000_029_038.txt")[:29000]
y = np.loadtxt("Modeling/y_train_30hrIF_30000_029_038.txt")[:29001,y_index]
yn_train = 0.1*(max(y) - min(y))*np.random.choice([-1,1],(y.shape))*np.random.random((y.shape))
y_n = y*(1+yn_train)

y_train = y[1:]
x_train = np.zeros((29000,2))
x_train[:,0] = u[:29000]
x_train[:,1] = y_n[:-1]

scaler_x = MinMaxScaler()
x_train = scaler_x.fit_transform(x_train)
scaler_y = MinMaxScaler()
y_train = scaler_y.fit_transform(y_train.reshape((y_train.shape[0],1)))

x_data, y_data = multivariate_data(x_train, y_train, 0, None, past, future,1)

OUTPUT_SIZE = 1
# %% Identifying the Model
model_1 = tf.keras.Sequential() 
model_1.add(tf.keras.layers.LSTM(5, input_shape=(past+future-1,2), activation = "sigmoid", return_sequences=True))
model_1.add(tf.keras.layers.LSTM(3, activation = "sigmoid"))
# model_1.add(tf.keras.layers.LSTM(3, activation = "sigmoid"))
# model_1.add(tf.keras.layers.Dense(5, input_shape=(past*2+future-1,), activation = "sigmoid"))  #, activation = 'sigmoid'
# model.add(Dense(50))
# model.add(tf.keras.layers.SimpleRNN(10, activation = 'tanh', return_sequences=True))
# model_1.add(tf.keras.layers.Dense(20, activation = 'sigmoid')) 
# model_1.add(tf.keras.layers.Dense(20, activation = 'sigmoid')) 
# model_1.add(tf.keras.layers.LSTM(20, return_sequences=True)) 
# model_1.add(tf.keras.layers.LSTM(5, activation = "sigmoid"))  #, return_sequences=True
# model_1.add(tf.keras.layers.Dense(20, activation = 'tanh'))  # , activation = 'sigmoid',activation="selu"
# model.add(tf.keras.layers.Dense(1, activation = 'selu'))
# add output layer
# model_2.add(tf.keras.layers.Dense(5, activation = 'tanh')) 
model_1.add(tf.keras.layers.Dense(OUTPUT_SIZE, activation = "sigmoid"))

opt = tf.keras.optimizers.Adam(learning_rate=0.005)
model_1.compile(optimizer=opt, loss='mse', metrics=['mean_squared_error'])

epoch = 500

callback = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=3)
Dense_history = model_1.fit(x_data,y_data.reshape((y_data.shape[0],y_data.shape[1])), epochs=epoch,
                               batch_size=200,
                                validation_split=0.15,
                           # steps_per_epoch = steps_per_epoch,
                                # validation_data = (x_t, y_t),
                                # validation_steps = validation_steps,
                               callbacks = [callback])

# %%
uv = np.loadtxt("Modeling/u_train_30hrIF_30000_029_038.txt")[29000:]
volMoisture = np.loadtxt("Modeling/y_train_30hrIF_30000_029_038.txt")[29000:,y_index]
a = 900
b = 1000
y_test = volMoisture[1:b+1]
x_test = np.zeros((b,2))
x_test[:,0] = uv[:b]
yn = 0.05*(max(y_test) - min(y_test))*np.random.choice([-1,1],(y_test.shape))*np.random.random((y_test.shape))
x_test[:,1] = volMoisture[:b] + yn

scaler_xt = MinMaxScaler()
x_test = scaler_xt.fit_transform(x_test)
y_test = y_test.reshape((y_test.shape[0],1))
scaler_yt = MinMaxScaler()
y_test = scaler_yt.fit_transform(y_test)


x_t, y_t = multivariate_data(x_test, y_test, 0, None, past, future,1)

# %%
y_pt_1 = model_3.predict(x_t)
ypt_1 = scaler_yt.inverse_transform(y_pt_1)
y_act = scaler_yt.inverse_transform(y_t.reshape((y_t.shape[0], y_t.shape[1])))

csfont = {'fontname':'Times New Roman'}
font = font_manager.FontProperties(family='Times New Roman',
                                    # weight='bold',
                                    style='normal', size=15)
plt.figure(3)
plt.plot(y_act[:,0], label = 'Actual Data')
plt.plot(ypt_1[:,0], label = 'Single-step-ahead Prediction')
# plt.plot(ypt_1[-500:,-1], label = 'Multi-step-ahead NN')
# plt.plot(ypt_2[30:,0], label = 'Multi-step-ahead NN -1st')
# plt.plot(ypt_2[:,-1], label = 'Multi-step-ahead NN - 30th')
plt.xlabel('Time Steps', **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)
plt.legend(frameon=False, ncol=1, prop=font)

# model_1.save("Modeling/p20_f1_5s_3s_sig_2h_029_038_4_35_e7_n")
# %% Multi-step ahead prediction performance

t_pred = 20  # 30 sampling times 
error = np.zeros(x_t.shape[0])
y_save = np.zeros((a,t_pred))

for t0 in range(a):
    print(t0)
    y_sim_1 = np.zeros((t_pred,1))
    NN0_1 = x_t[t0,:]
    # y_sub = np.array([DM(M1(NN0_1)), DM(M2(NN0_1)), DM(M3(NN0_1))])  #, DM(M4(NN0_1)), DM(M5(NN0_1))])
    y_sim_1[0,:] = np.array(DM(M3(NN0_1))).flatten()
    u_sim = x_t[t0:t0+t_pred,-1,0]
    for i in range(t_pred-1):
        NN0_1 = np.concatenate((NN0_1[1:], np.array([u_sim[i], y_sim_1[i,0]]).reshape((1,2))))
        # y_sub = np.array([DM(M1(NN0_1)), DM(M2(NN0_1)), DM(M3(NN0_1))])  #, DM(M4(NN0_1)), DM(M5(NN0_1))])
        y_sim_1[i+1,:] = np.array(DM(M3(NN0_1))).flatten()
    yp_1 = scaler_yt.inverse_transform(y_sim_1).flatten()
    y_save[t0,:] = yp_1
    error[t0] = sum((y_act[t0:t0+t_pred].flatten() - yp_1)**2/y_act[t0:t0+t_pred].flatten())
# yp_2 = scaler_yt.inverse_transform(y_sim_2)
# yp = (yp_1 + yp_2)/2
# np.savetxt("Modeling/y_pred_multi_20hrIF_1000_029_038.txt", y_save)
# %%
# y_save = np.loadtxt("y_pred_multi_20hrIF_1000_029_038.txt")
a = 900
csfont = {'fontname':'Times New Roman'}
font = font_manager.FontProperties(family='Times New Roman',
                                   # weight='bold',
                                    style='normal', size=17)
tplot = np.arange(a)
b = 250
c = 370

fig,ax1 = plt.subplots()
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
ax1.plot(tplot, y_act[:a], label = 'Actual Trajectory', color = "k", linewidth = 3)
ax1.plot(tplot, y_save[:,0], label = "Single-step-ahead") # 0.000914381408118862
# ax1.plot(tplot[:], y_save[:,4], label = "5-step-ahead", color = "r")  # 0.0017256467027993415
ax1.plot(tplot[9:], y_save[:a-9,9], label = "Ten-step-ahead", color = "r", linestyle = "dashdot") # 0.002444138910936656
ax1.plot(tplot[19:], y_save[:a-19,19], label = "Twenty-step-ahead", color = "y", linestyle = "dashed") # 0.0030556959486465352
ax1.legend(frameon=False, bbox_to_anchor=(0.1,1,1,0), ncol=2, prop=font)
plt.xlabel("Time Steps", **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)


# ax2 = plt.axes([0,0,1,1])
# ip = InsetPosition(ax1, [0.2,0.13,0.2,0.2])
# ax2.set_axes_locator(ip)
# mark_inset(ax1, ax2, loc1=2, loc2=4, fc="none", ec='0.6')
# ax2.plot(tplot[b:c+20], y_act[b:c+20], label = 'Actual Trajectory', color = "k", linewidth = 3)
# ax2.plot(tplot[b:c], y_save[b:c,0], label = "Single-step-ahead") 
# ax2.plot(tplot[b:c+9], y_save[b-9:c,9], label = "Ten-step-ahead", color = "r", linestyle = "dashdot") 
# ax2.plot(tplot[b:c+19], y_save[b-19:c,19], label = "Twenty-step-ahead", color = "y", linestyle = "dashed") 
# # ax2.legend(loc=0)
# # ax2.set_xticklabels(ax2.get_xticks(), backgroundcolor='w')
# ax2.tick_params(axis='x', which='major', pad=8)
# plt.suptitle('Composite Plot with original subsystems(Ca)', fontsize=15)
plt.show()

# tplot = np.arange(a)


# plt.figure(4)
# plt.plot(tplot, y_act[:a], label = 'Actual Trajectory', color = "k", linewidth = 4)
# plt.plot(tplot, y_save[:,0], label = "1-step-ahead", color = "r") # 0.0017196699344804148
# # plt.plot(tplot[4:], y_save[:-4,4], label = "5-step-ahead", color = "r")  # 0.003124494079695719
# # plt.plot(tplot[:], y_save[:,9], label = "10-step-ahead", color = "y") #0.00415294297808327
# plt.plot(tplot[:], y_save[:,19], label = "20-step-ahead", color = "g") # 0.004423033515746692
# plt.legend(frameon=False, loc = "upper center", ncol=1, prop=font) # bbox_to_anchor=(0.3,0.2,0.3,0.3)
# plt.xlabel("Time Steps", **csfont, fontsize = 17)
# plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)

e_0 = np.sqrt(sum((y_act[:900,0] - y_save[:,0])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0])) #(max((y_act[:900,:] - y_save[:,:]).flatten()))
e_5 = np.sqrt(sum((y_act[4:904,0] - y_save[:,4])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0])) 
e_10 = np.sqrt(sum((y_act[9:909,0] - y_save[:,9])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0]))
e_20 = np.sqrt(sum((y_act[19:919,0] - y_save[:,19])**2)/900)/(max(y_act[:900,0]) - min(y_act[:900,0]))

print(e_0*100)
print(e_5*100)
print(e_10*100)
print(e_20*100)

# e_0 = np.sqrt(sum((y_act[:900,0] - y_save[:,0])**2/y_act[:900,0]**2))/900
# e_5 = np.sqrt(sum((y_act[4:904,0] - y_save[:,4])**2/y_act[4:904,0]**2))/900
# e_10 = np.sqrt(sum((y_act[9:909,0] - y_save[:,9])**2/y_act[9:909,0]**2))/900
# e_20 = np.sqrt(sum((y_act[19:919,0] - y_save[:,19])**2/y_act[19:919,0]**2))/900

# np.savetxt("u_vali_20hrIF_1000_029_038.txt", uv)
# np.savetxt("x_vali_20hrIF_1000_029_038.txt", headArray)
# np.savetxt("y_vali_20hrIF_1000_029_038.txt", volMoisture)
# np.savetxt("y_pred_multi_20hrIF_1000_029_038.txt", y_save)
# t0=200
# plt.figure(4)
# plt.plot(y_act[t0:t0+t_pred,-1], label = 'Actual Data', color = "k")  # , marker = 'o'
# plt.plot(ypt_1[t0:t0+t_pred,-1], label = 'Single-step', color = "b")  # , marker = 'd'
# plt.plot(yp_1[0:,-1], label = 'Multi-step', color = "r")  # , marker = '^'
# plt.xlabel('Time Steps', **csfont, fontsize = 17)
# plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)

# plt.legend(frameon=False, ncol=1, prop=font)

